"""
author:dlr123
date:2021年09月28日
"""
from collections import defaultdict

import numpy as np
import pandas as pd

from tree.tree import TreeNode
from opt_func.FuncApi import Function
import time


class MCTS:

    def __init__(self,
                 policy_value_fn: Function,
                 c_values: float = 1.0,
                 n_playout: int = 10000,
                 simu_times: int = 100,
                 pause_times: int = 50,
                 pause_thre: float = 0.02,
                 stop_layer_play_time: int = 100,
                 extra_layer: int = -1):
        """
        初始化MCTS
        :param policy_value_fn: 目标函数，需要是Function类型
        :param c_values: uts分数中的超参数，用以调整控制广度和深度的关系，越小越倾向于深度优先
        :param n_playout: 模型训练的上限
        :param simu_times: 每次仿真的次数
        :param pause_times: 当搜寻到最优值pause_times不变时停止
        :param pause_thre: 用以判断当前最优值是否不变
        :param stop_layer_play_time: 待提取层中最低层至少需要训练的次数
        :param extra_layer: 需要提取的层数
        """

        self.last_best_list = []
        self.starttime = 0
        self.extra_layer = extra_layer
        self.stop_layer_play_time = stop_layer_play_time

        self.pause_thre = pause_thre
        self.pause_times = pause_times
        self.invariant_times = 0

        self.root = TreeNode()
        self.policy_value_fn = policy_value_fn
        values_range = policy_value_fn.values_range
        # 如果输入的是字典型的变量取值范围，可以方便多树的生成
        if isinstance(values_range, dict):
            self.values_range_dict = values_range
            self.values_range_list = list(self.values_range_dict.values())
            self.values_set = list(self.values_range_dict.keys())
            self.values_items = dict(
                zip(self.values_set, range(len(self.values_set))))
        else:
            self.values_range_list = values_range
        self.value_index = policy_value_fn.data_index
        self.value_num = len(self.values_range_list)  # value_num：表示变量的数量
        self.c_values = c_values
        self.n_playout = n_playout
        self.simu_times = simu_times
        self.json_format = None
        self.visit_times = defaultdict(int)
        self.time = 0

    def playout(self, select_rate=0.5):
        """
        :param select_rate: 选择时倾向于探索未知结点的
        :return:
        """
        node = self.root
        values_list = []
        temp_node = node
        values = None
        # 找寻本轮进行仿真的叶节点
        while True:
            if node.is_leaf():
                break
            if isinstance(select_rate, float):
                values, temp_node = node.select(self.c_values, set(
                    self.values_range_list[node.depth + 1]), thre=select_rate)
            # 如果当前选择的节点是并未访问过的节点，则对其进行访问，赋予其初始值
            else:
                node_depth = node.depth + 1
                if len(select_rate) >= (node_depth + 1):
                    values, temp_node = node.select(self.c_values, set(self.values_range_list[node.depth + 1]),
                                                    thre=select_rate[node_depth])
                else:
                    values, temp_node = node.select(self.c_values, set(self.values_range_list[node.depth + 1]),
                                                    thre=0.5)
            if temp_node is None:
                break
            # 如果当前节点既不是叶节点也曽被访问过，则继续迭代
            self.visit_times[node.depth] += 1
            node = temp_node
            # 记录迭代过程中的变量路径
            values_list.append(values)
        # 对叶节点进行更新
        self.leaf_expand_update(node, values_list, temp_node, values)

    def leaf_expand_update(self, node, values_list, temp_node=None, temp_value=None):
        """:cvar
        考虑当前节点是叶节点还是未更新节点
        考虑当前固定的变量数是否完全，未完全时需要进行仿真
        """
        # 首先判断当前待更新节点是叶节点还是未访问节点
        if temp_node is None:
            values_list.append(temp_value)
        else:
            if node.depth != self.value_num - 1:
                temp_value = np.random.choice(
                    self.values_range_list[node.depth + 1])
                values_list.append(temp_value)
        # 当变量列表依据固定完全时，直接计算，不需要进行仿真
        if len(values_list) == self.value_num:
            policy_value = self.policy_value_fn(values_list, None)
            if node.depth == self.value_num - 1:
                node.update_recursive(policy_value)
                self.visit_times[node.depth] += 1  # 访问次数加1
                return
            else:
                node.expand(temp_value, policy_value)
        else:
            # 当变量没有被完全决定完时需要进行仿真更新
            policy_value = self._mc_simulink(node.depth + 1, values_list)
            node.expand(temp_value, policy_value)
        self.visit_times[node.depth] += 1
        self.visit_times[node.depth + 1] += 1  # 层级访问次数加1

    def _mc_simulink(self, depth, value_list):
        """:var
        进行蒙特卡洛仿真
        label表示是第几个变量
        """
        # TODO:设计违反约束的惩罚函数
        # TODO:设计判断是否违反约束的方法

        def fn(res_value): return self.policy_value_fn(value_list, res_values)
        res_values = np.array([np.random.choice(ls, self.simu_times)
                               for ls in self.values_range_list[(depth + 1):]])
        result = fn(res_values)
        return result

    def run(self, select_rate=0.5):
        """:cvar
        模型运行函数
        """
        self.starttime = time.time()
        for i in range(self.n_playout):
            if (i + 1) % 5000 == 0:
                print(f"孩子已经学了{i + 1}轮啦！")
            last_value = self.root.min_score
            self.playout(select_rate=select_rate)
            is_pause = self.is_pause(last_value)
            if self.extra_layer != -1:
                if(self.visit_times[self.extra_layer - 1] >= self.stop_layer_play_time) and is_pause:
                    break
        print(f"在训练{i - self.invariant_times + 1}轮后即可结束",
              f"共调用了{self.simu_times * (i - self.invariant_times + 1)}次目标函数",
              f"总共执行了:{self.time}s", sep='\n')

    def is_pause(self, last_value):
        flag1 = False  # 判断最优值是否有变化
        flag2 = True  # 判断最优变量组合是否有变化
        change_rate = (last_value - self.root.min_score)/(last_value + 1e-8)
        if change_rate <= self.pause_thre:
            flag1 = True
        if (self.invariant_times % 100) == 0:
            temp_list = self.output_the_best(is_use2pause=True)
            if self.last_best_list != temp_list:
                flag2 = False
                self.last_best_list = temp_list
        if flag1 and flag2:
            self.invariant_times += 1
        else:
            self.invariant_times = 0
            self.time = time.time() - self.starttime
        if self.invariant_times == self.pause_times:
            return True
        return False

    def output_the_best(self, is_use2pause=False):
        """:var
        输出最优结果
        """
        value_list = []
        node = self.root
        while not node.is_leaf():
            value, node = node.select(self.c_values, is_output=True)
            value_list.append(value)
        if not is_use2pause:
            if len(value_list) != self.value_num:
                print("树的最优解并未完全展开")
                print(f"当前还有{self.value_num - len(value_list)}层没有展开")
            for i, values in enumerate(value_list):
                print(f"第{self.value_index[i] + 1}个变量取值:{values}")
            print(f"最终目标值为:{node.min_score}")
            return
        return value_list

    # 接下来的模块是用于结果的数据结构输出的
    def selecct_children_node(self, parent, threshold=0.01):
        """:cvar
        当前版本的实现方法中，在训练过程中所有节点都是访问过的，所以不需要第一层过滤
        第二层过滤是根据目标分数的阈值进行过滤的，同时考虑了得分和方差置信度
        """
        # TODO:考虑如何过滤更加合理，用什么样的参数或指标进行过滤
        stay = {}
        if not parent.is_leaf():
            children = pd.Series(parent.children)
            # children = children[children.apply(lambda node : node.n_visits) != 0]  # 过滤掉没有被访问过的节点
            # children_score = children.apply(lambda node : node.get_values(self.c_values))  # 映射出访问过的节点各自的最小值
            children_score = children.apply(
                lambda node: node.min_score)  # 映射出访问过的节点各自的最小值
            threshold_score = children_score.quantile(threshold)  # 设置阈值门槛
            stay = children[children_score <=
                            threshold_score].to_dict()  # 过滤后剩余的节点
        return stay

    def _recurve_output(self, CurNode, thre=1, format='full', is_show_min=False, is_show_info=False, stop_layer=-1):
        Curchild = self.selecct_children_node(CurNode, threshold=thre)
        if Curchild != {}:
            child_dict = {}
            for child_id in Curchild:
                child = Curchild[child_id]

                # 设置输出停止的层数
                if child.depth == stop_layer:
                    return {}
                output = self._recurve_output(child, thre=thre, format=format,
                                              is_show_min=is_show_min, is_show_info=is_show_info, stop_layer=stop_layer)
                if format == 'full':
                    if is_show_info:
                        info = {'visit_t': child.n_visits,
                                'min_score': child.min_score, 'label': child.depth}
                        # child_dict[self.values_range_list[child.depth][child_id]] = [info,output]
                        child_dict[child_id] = [info, output]
                    else:
                        # child_dict[self.values_range_list[child.depth][child_id]] = [child,output]
                        child_dict[child_id] = [child, output]
                elif format == 'simple':
                    if is_show_min:
                        # child_dict[self.values_range_list[child.depth][child_id]] = [child.min_score,output]
                        child_dict[child_id] = [child.min_score, output]
                    else:
                        # child_dict[self.values_range_list[child.depth][child_id]] = output
                        child_dict[child_id] = output
            return child_dict
        return {}

    def output_tree(self, thre=1, format='full', is_show_min=False, is_show_info=False):
        node = self.root
        return self._recurve_output(node, thre=thre, format=format,
                                    is_show_min=is_show_min, is_show_info=is_show_info, stop_layer=self.extra_layer)

    def merge_layer(self, thre=1):
        """
        :param thre: 设置每层提取时的阈值
        :return:
        """

        # TODO:对各层设计一个评价指标
        simple_tree = self.output_tree(format='simple', thre=thre)
        all_layers = []
        cur_layer = [simple_tree]
        while cur_layer != []:
            next_layer = []
            cur_values = set()
            for node in cur_layer:
                cur_values.update(list(node.keys()))
                next_layer.extend(list(node.values()))
                while {} in next_layer:
                    next_layer.remove({})
            cur_layer = next_layer
            all_layers.append(list(cur_values))
        return all_layers
